"""
Cloudflare R2 uploader using S3-compatible API.

Reference: https://developers.cloudflare.com/r2/api/s3/api/

Capabilities
- Upload from file path, bytes, or base64 (data URL supported)
- Path-style addressing for R2 (bucket in path)
- Optional public URL base for returned link (e.g., custom domain or r2.dev domain)
- Minimal, explicit return structure for easy plumbing

Dependencies
- boto3 + botocore
  Install: pip install boto3 botocore

Environment (optional)
- R2_ACCOUNT_ID, R2_ACCESS_KEY_ID, R2_SECRET_ACCESS_KEY, R2_BUCKET
- R2_PUBLIC_BASE_URL (e.g., https://pub-xxxxxxxxxxxx.r2.dev/your-bucket)

CLI usage
  python agent/cf_uploader.py --account-id <id> --access-key-id <ak> --secret-access-key <sk> \
      --bucket my-bucket --file path/to/img.png --key images/img.png --public-base-url https://pub-xxx.r2.dev/my-bucket

  python agent/cf_uploader.py --bucket my-bucket --b64 "data:image/png;base64,..." --name img.png \
      --account-id $R2_ACCOUNT_ID --access-key-id $R2_ACCESS_KEY_ID --secret-access-key $R2_SECRET_ACCESS_KEY
"""

from __future__ import annotations

import base64
import mimetypes
import os
from dataclasses import dataclass
from typing import Any, Dict, Optional


def _strip_data_url_prefix(s: str) -> str:
    if isinstance(s, str) and s.startswith("data:"):
        i = s.find(",")
        if i != -1:
            return s[i + 1 :]
    return s


def _guess_mime(name: str) -> str:
    mt, _ = mimetypes.guess_type(name)
    return mt or "application/octet-stream"


@dataclass
class UploadResult:
    ok: bool
    status_code: int
    key: Optional[str]
    url: Optional[str]
    etag: Optional[str]
    version_id: Optional[str]
    error: Optional[str] = None
    extra: Optional[Dict[str, Any]] = None


class CfR2Uploader:
    def __init__(
        self,
        account_id: str,
        access_key_id: str,
        secret_access_key: str,
        bucket: str,
        region: str = "auto",
        endpoint_url: Optional[str] = None,
        public_base_url: Optional[str] = None,
    ) -> None:
        try:
            import boto3  # type: ignore
            from botocore.config import Config  # type: ignore
        except Exception as e:
            raise RuntimeError(
                "boto3/botocore is required. Install with: pip install boto3 botocore"
            ) from e

        self.bucket = bucket
        self.public_base_url = public_base_url.rstrip("/") if public_base_url else None
        self.endpoint_url = endpoint_url or f"https://{account_id}.r2.cloudflarestorage.com"

        cfg = Config(s3={"addressing_style": "path"})
        self._s3 = boto3.client(
            "s3",
            endpoint_url=self.endpoint_url,
            aws_access_key_id=access_key_id,
            aws_secret_access_key=secret_access_key,
            region_name=region,
            config=cfg,
        )

    # ----------------------------
    # Core
    # ----------------------------
    def _public_url(self, key: str) -> Optional[str]:
        if self.public_base_url:
            return f"{self.public_base_url}/{key.lstrip('/')}"
        # Fallback to API endpoint URL (may not be publicly accessible)
        return f"{self.endpoint_url}/{self.bucket}/{key.lstrip('/')}"

    def _put_object(self, key: str, body: bytes, content_type: Optional[str] = None, extra_args: Optional[Dict[str, Any]] = None) -> UploadResult:
        try:
            kwargs: Dict[str, Any] = {
                "Bucket": self.bucket,
                "Key": key,
                "Body": body,
            }
            if content_type:
                kwargs["ContentType"] = content_type
            if extra_args:
                kwargs.update(extra_args)
            resp = self._s3.put_object(**kwargs)
            etag = resp.get("ETag")
            version_id = resp.get("VersionId")
            return UploadResult(
                ok=True,
                status_code=200,
                key=key,
                url=self._public_url(key),
                etag=etag.strip('"') if isinstance(etag, str) else None,
                version_id=version_id if isinstance(version_id, str) else None,
                extra={"raw": resp},
            )
        except Exception as e:
            # Try to pull status code from exception
            sc = getattr(e, "response", {}).get("ResponseMetadata", {}).get("HTTPStatusCode", 0) if hasattr(e, "response") else 0
            return UploadResult(ok=False, status_code=int(sc or 0), key=key, url=None, etag=None, version_id=None, error=str(e))

    # ----------------------------
    # Public helpers
    # ----------------------------
    def upload_bytes(self, name: str, data: bytes, key: Optional[str] = None, content_type: Optional[str] = None) -> UploadResult:
        if not isinstance(data, (bytes, bytearray)):
            return UploadResult(ok=False, status_code=0, key=None, url=None, etag=None, version_id=None, error="data must be bytes")
        ct = content_type or _guess_mime(name)
        obj_key = key or name
        return self._put_object(obj_key, bytes(data), content_type=ct)

    def upload_base64(self, name: str, b64: str, key: Optional[str] = None, content_type: Optional[str] = None) -> UploadResult:
        b64 = _strip_data_url_prefix(b64 or "")
        try:
            raw = base64.b64decode(b64, validate=False)
        except Exception as e:
            return UploadResult(ok=False, status_code=0, key=None, url=None, etag=None, version_id=None, error=f"invalid base64: {e}")
        return self.upload_bytes(name=name, data=raw, key=key, content_type=content_type)

    def upload_file(self, path: str, key: Optional[str] = None, content_type: Optional[str] = None) -> UploadResult:
        if not os.path.isfile(path):
            return UploadResult(ok=False, status_code=0, key=None, url=None, etag=None, version_id=None, error=f"file not found: {path}")
        name = os.path.basename(path)
        ct = content_type or _guess_mime(name)
        obj_key = key or name
        try:
            with open(path, "rb") as f:
                data = f.read()
        except Exception as e:
            return UploadResult(ok=False, status_code=0, key=None, url=None, etag=None, version_id=None, error=str(e))
        return self._put_object(obj_key, data, content_type=ct)


if __name__ == "__main__":
    import argparse
    import json as _json

    p = argparse.ArgumentParser(description="Upload to Cloudflare R2 (S3 API)")
    p.add_argument("--account-id", default=os.environ.get("R2_ACCOUNT_ID"))
    p.add_argument("--access-key-id", default=os.environ.get("R2_ACCESS_KEY_ID"))
    p.add_argument("--secret-access-key", default=os.environ.get("R2_SECRET_ACCESS_KEY"))
    p.add_argument("--bucket", default=os.environ.get("R2_BUCKET"))
    p.add_argument("--endpoint-url", default=None, help="Override endpoint, default https://<account>.r2.cloudflarestorage.com")
    p.add_argument("--public-base-url", default=os.environ.get("R2_PUBLIC_BASE_URL"))
    p.add_argument("--region", default="auto")

    g = p.add_mutually_exclusive_group(required=True)
    g.add_argument("--file", help="Path to file")
    g.add_argument("--b64", help="Base64 string (optionally data URL)")
    p.add_argument("--name", help="Filename to use when --b64 is provided")
    p.add_argument("--key", help="Object key in bucket (defaults to filename)")
    args = p.parse_args()

    missing = [k for k, v in {
        "account-id": args.account_id,
        "access-key-id": args.access_key_id,
        "secret-access-key": args.secret_access_key,
        "bucket": args.bucket,
    }.items() if not v]
    if missing:
        raise SystemExit("Missing required params: " + ", ".join(missing))

    uploader = CfR2Uploader(
        account_id=args.account_id,
        access_key_id=args.access_key_id,
        secret_access_key=args.secret_access_key,
        bucket=args.bucket,
        region=args.region,
        endpoint_url=args.endpoint_url,
        public_base_url=args.public_base_url,
    )

    if args.file:
        res = uploader.upload_file(args.file, key=args.key)
    else:
        if not args.name:
            raise SystemExit("--name is required when using --b64")
        res = uploader.upload_base64(args.name, args.b64, key=args.key)

    print(_json.dumps({
        "ok": res.ok,
        "status_code": res.status_code,
        "key": res.key,
        "url": res.url,
        "etag": res.etag,
        "version_id": res.version_id,
        "error": res.error,
    }, ensure_ascii=False, indent=2))

